package shardctrler

import (
	"sort"
	"sync"
	"time"

	"6.824/labgob"
	"6.824/labrpc"
	"6.824/raft"
)

type ShardCtrler struct {
	mu      sync.Mutex
	me      int
	rf      *raft.Raft
	applyCh chan raft.ApplyMsg

	// Your data here.

	configs []Config // indexed by config num

	applyIndex int //persist
	fromIndex  int

	identifySeqs  map[int64]int64 //persist
	replyChannels []chan int
}

type Op struct {
	// Your data here.
	OpType      string
	Servers     map[int][]string
	GIDs        []int
	Shard       int
	GID         int
	ClinetId    int64
	IdentifySeq int64
}

func (sc *ShardCtrler) waitLogApply(index int, term int) Err {
	sc.mu.Lock()
	if len(sc.replyChannels) > sc.adjustIndex(index) {
		sc.replyChannels[sc.adjustIndex(index)] = make(chan int, 1)
	} else {
		for len(sc.replyChannels)-1 < sc.adjustIndex(index-1) {
			sc.replyChannels = append(sc.replyChannels, nil)
		}
		sc.replyChannels = append(sc.replyChannels, make(chan int, 1))
	}

	toReadChannel := sc.replyChannels[sc.adjustIndex(index)]
	sc.mu.Unlock()

	select {
	case readTerm := <-toReadChannel:

		if readTerm != term {
			toReadChannel <- readTerm // give back
			return Err("apply failed")
		}

		// append success
		return OK

	case <-time.After(2000 * time.Millisecond):
		// timeout: 2000ms
		return Err("timeout")
	}
}

func (sc *ShardCtrler) adjustIndex(index int) int {
	return index - sc.fromIndex
}

func (sc *ShardCtrler) Join(args *JoinArgs, reply *JoinReply) {
	// Your code here.
	sc.mu.Lock()
	send_args := Op{"Join", args.Servers, nil, -1, -1, args.ClinetId, args.IdentifySeq}
	index, term, isLeader := sc.rf.Start(send_args)
	if !isLeader {
		sc.mu.Unlock()
		reply.WrongLeader = true
		reply.Err = Err("wrong guy")
		return
	}
	sc.mu.Unlock()
	reply.WrongLeader = false

	applyRes := sc.waitLogApply(index, term)
	reply.Err = applyRes

}

func (sc *ShardCtrler) Leave(args *LeaveArgs, reply *LeaveReply) {
	// Your code here.

	sc.mu.Lock()
	send_args := Op{"Leave", nil, args.GIDs, -1, -1, args.ClinetId, args.IdentifySeq}
	index, term, isLeader := sc.rf.Start(send_args)
	if !isLeader {
		sc.mu.Unlock()
		reply.WrongLeader = true
		reply.Err = Err("wrong guy")
		return
	}
	sc.mu.Unlock()
	reply.WrongLeader = false

	applyRes := sc.waitLogApply(index, term)
	reply.Err = applyRes

}

func (sc *ShardCtrler) Move(args *MoveArgs, reply *MoveReply) {
	// Your code here.

	sc.mu.Lock()
	send_args := Op{"Move", nil, nil, args.Shard, args.GID, args.ClinetId, args.IdentifySeq}
	index, term, isLeader := sc.rf.Start(send_args)
	if !isLeader {
		sc.mu.Unlock()
		reply.WrongLeader = true
		reply.Err = Err("wrong guy")
		return
	}
	sc.mu.Unlock()
	reply.WrongLeader = false

	applyRes := sc.waitLogApply(index, term)
	reply.Err = applyRes

}

func (sc *ShardCtrler) Query(args *QueryArgs, reply *QueryReply) {
	// Your code here.

	sc.mu.Lock()
	send_args := Op{"Query", nil, nil, -1, -1, args.ClinetId, args.IdentifySeq}
	index, term, isLeader := sc.rf.Start(send_args)
	if !isLeader {
		sc.mu.Unlock()
		reply.WrongLeader = true
		reply.Err = Err("wrong guy")
		return
	}
	sc.mu.Unlock()
	reply.WrongLeader = false

	applyRes := sc.waitLogApply(index, term)

	if applyRes == OK {
		reply.Err = OK
		sc.mu.Lock()
		if args.Num >= 0 && args.Num < len(sc.configs) {
			reply.Config = sc.configs[args.Num]
		} else if args.Num == -1 || args.Num >= len(sc.configs) {
			reply.Config = sc.configs[len(sc.configs)-1]
		} else {
			raft.Assert(false, "Unknown index")
		}
		sc.mu.Unlock()
	} else {
		reply.Err = applyRes
	}

}

func mapKeys(m map[int][]string) []int {
	ret := []int{}
	for key, _ := range m {
		ret = append(ret, key)
	}
	sort.Slice(ret, func(i, j int) bool { return ret[i] < ret[j] })
	return ret
}

func (sc *ShardCtrler) selectApplyLogs(opArgs Op, term int) {
	if sc.identifySeqs[opArgs.ClinetId] == opArgs.IdentifySeq {
		sc.applyIndex++
		if len(sc.replyChannels) > sc.adjustIndex(sc.applyIndex) && sc.replyChannels[sc.adjustIndex(sc.applyIndex)] != nil {
			sc.replyChannels[sc.adjustIndex(sc.applyIndex)] <- term
		}
		return
	}

	lastConfig := sc.configs[len(sc.configs)-1]

	switch opArgs.OpType {
	case "Join":
		updatedGroups := make(map[int][]string)
		for key, value := range lastConfig.Groups {
			updatedGroups[key] = value
		}

		for key, value := range opArgs.Servers {
			updatedGroups[key] = value
		}

		numGroups := len(updatedGroups)
		keys := mapKeys(updatedGroups)
		updatedShards := [NShards]int{}
		for i := 0; i < NShards; i++ {
			if numGroups != 0 {
				updatedShards[i] = keys[i%numGroups]
			} else {
				updatedShards[i] = 0
			}
		}

		insertConfig := Config{lastConfig.Num + 1, updatedShards, updatedGroups}
		sc.configs = append(sc.configs, insertConfig)

	case "Leave":
		updatedGroups := make(map[int][]string)
		for key, value := range lastConfig.Groups {
			updatedGroups[key] = value
		}

		for _, value := range opArgs.GIDs {
			delete(updatedGroups, value)
		}

		numGroups := len(updatedGroups)
		keys := mapKeys(updatedGroups)
		updatedShards := [NShards]int{}
		for i := 0; i < NShards; i++ {
			if numGroups != 0 {
				updatedShards[i] = keys[i%numGroups]
			} else {
				updatedShards[i] = 0
			}
		}

		insertConfig := Config{lastConfig.Num + 1, updatedShards, updatedGroups}
		sc.configs = append(sc.configs, insertConfig)

	case "Move":
		updatedGroups := make(map[int][]string)
		for key, value := range lastConfig.Groups {
			updatedGroups[key] = value
		}

		updatedShards := lastConfig.Shards
		raft.Assert(&updatedShards != &lastConfig.Shards, "not &updatedShards!=&lastConfig.Shards")

		updatedShards[opArgs.Shard] = opArgs.GID

		insertConfig := Config{lastConfig.Num + 1, updatedShards, updatedGroups}
		sc.configs = append(sc.configs, insertConfig)

	case "Query":
		// nothing to do
	default:
		panic("Unknown OP")
	}

	sc.identifySeqs[opArgs.ClinetId] = opArgs.IdentifySeq
	sc.applyIndex++
	if len(sc.replyChannels) > sc.adjustIndex(sc.applyIndex) && sc.replyChannels[sc.adjustIndex(sc.applyIndex)] != nil {
		sc.replyChannels[sc.adjustIndex(sc.applyIndex)] <- term
	}
}

func (sc *ShardCtrler) applyLogDaemon() {
	for msg := range sc.applyCh {
		if msg.SnapshotValid {
			raft.Assert(false, "should not reach here")
			continue
		}

		raft.Assert(msg.CommandIndex == sc.applyIndex+1, "not msg.CommandIndex==kv.applyIndex+1")

		cmd := msg.Command.(Op)
		term := msg.CommandTerm

		sc.mu.Lock()
		sc.selectApplyLogs(cmd, term)
		sc.mu.Unlock()

	}
}

// the tester calls Kill() when a ShardCtrler instance won't
// be needed again. you are not required to do anything
// in Kill(), but it might be convenient to (for example)
// turn off debug output from this instance.
func (sc *ShardCtrler) Kill() {
	sc.rf.Kill()
	// Your code here, if desired.
}

// needed by shardkv tester
func (sc *ShardCtrler) Raft() *raft.Raft {
	return sc.rf
}

// servers[] contains the ports of the set of
// servers that will cooperate via Raft to
// form the fault-tolerant shardctrler service.
// me is the index of the current server in servers[].
func StartServer(servers []*labrpc.ClientEnd, me int, persister *raft.Persister) *ShardCtrler {
	sc := new(ShardCtrler)
	sc.me = me

	sc.configs = make([]Config, 1)
	sc.configs[0].Groups = map[int][]string{}

	labgob.Register(Op{})
	sc.applyCh = make(chan raft.ApplyMsg)
	sc.rf = raft.Make(servers, me, persister, sc.applyCh)

	// Your code here.
	sc.configs = []Config{{0, [NShards]int{}, make(map[int][]string)}}
	sc.identifySeqs = make(map[int64]int64)

	go sc.applyLogDaemon()

	return sc
}
